#ifndef _SERIALIZABLEMAP_H
#define _SERIALIZABLEMAP_H

#include "_global.h"
#include "jansson/jansson_x64dbg.h"
#include "memory.h"
#include "module.h"
#include "threading.h"

template <class TValue>
class JSONWrapper {
 public:
  virtual ~JSONWrapper() {}

  void SetJson(JSON json) { mJson = json; }

  virtual bool Save(const TValue& value) = 0;
  virtual bool Load(TValue& value) = 0;

 protected:
  void setString(const char* key, const std::string& value) {
    set(key, json_string(value.c_str()));
  }

  bool getString(const char* key, std::string& dest) const {
    auto jsonValue = get(key);
    if (!jsonValue) return false;
    auto str = json_string_value(jsonValue);
    if (!str) return false;
    dest = str;
    return true;
  }

  void setHex(const char* key, duint value) { set(key, json_hex(value)); }

  bool getHex(const char* key, duint& value) const {
    auto jsonValue = get(key);
    if (!jsonValue) return false;
    value = duint(json_hex_value(jsonValue));
    return true;
  }

  void setBool(const char* key, bool value) { set(key, json_boolean(value)); }

  bool getBool(const char* key, bool& value) const {
    auto jsonValue = get(key);
    if (!jsonValue) return false;
    value = json_boolean_value(jsonValue);
    return true;
  }

  template <typename T>
  void setInt(const char* key, T value) {
    set(key, json_integer(value));
  }

  template <typename T>
  bool getInt(const char* key, T& value) {
    auto jsonValue = get(key);
    if (!jsonValue) return false;
    value = T(json_integer_value(jsonValue));
    return true;
  }

  // ReSharper disable once CppMemberFunctionMayBeConst
  void set(const char* key, JSON value) {
    json_object_set_new(mJson, key, value);
  }

  JSON get(const char* key) const { return json_object_get(mJson, key); }

  JSON mJson = nullptr;
};

template <SectionLock TLock, class TKey, class TValue, class TMap,
          class TSerializer>
class SerializableTMap {
  static_assert(std::is_base_of<JSONWrapper<TValue>, TSerializer>::value,
                "TSerializer is not derived from JSONWrapper<TValue>");

 public:
  using TValuePred = std::function<bool(const TValue& value)>;

  virtual ~SerializableTMap() {}

  bool Add(const TValue& value) {
    EXCLUSIVE_ACQUIRE(TLock);
    return addNoLock(value);
  }

  bool Get(const TKey& key, TValue& value) const {
    SHARED_ACQUIRE(TLock);
    auto found = mMap.find(key);
    if (found == mMap.end()) return false;
    value = found->second;
    return true;
  }

  bool Contains(const TKey& key) const {
    SHARED_ACQUIRE(TLock);
    return mMap.count(key) > 0;
  }

  bool Delete(const TKey& key) {
    EXCLUSIVE_ACQUIRE(TLock);
    return mMap.erase(key) > 0;
  }

  void DeleteWhere(TValuePred predicate) {
    EXCLUSIVE_ACQUIRE(TLock);
    for (auto itr = mMap.begin(); itr != mMap.end();) {
      if (predicate(itr->second))
        itr = mMap.erase(itr);
      else
        ++itr;
    }
  }

  bool GetWhere(TValuePred predicate, TValue& value) {
    return getWhere(predicate, &value);
  }

  bool GetWhere(TValuePred predicate) { return getWhere(predicate, nullptr); }

  void Clear() {
    EXCLUSIVE_ACQUIRE(TLock);
    TMap empty;
    std::swap(mMap, empty);
  }

  void CacheSave(JSON root) const {
    SHARED_ACQUIRE(TLock);
    auto jsonValues = json_array();
    TSerializer serializer;
    for (const auto& itr : mMap) {
      auto jsonValue = json_object();
      serializer.SetJson(jsonValue);
      if (serializer.Save(itr.second)) json_array_append(jsonValues, jsonValue);
      json_decref(jsonValue);
    }
    if (json_array_size(jsonValues))
      json_object_set(root, jsonKey(), jsonValues);
    json_decref(jsonValues);
  }

  void CacheLoad(JSON root, const char* keyprefix = nullptr) {
    EXCLUSIVE_ACQUIRE(TLock);
    auto jsonValues = json_object_get(
        root, keyprefix ? (keyprefix + String(jsonKey())).c_str() : jsonKey());
    if (!jsonValues) return;
    size_t i;
    JSON jsonValue;
    TSerializer deserializer;
    json_array_foreach(jsonValues, i, jsonValue) {
      deserializer.SetJson(jsonValue);
      TValue value;
      if (deserializer.Load(value)) addNoLock(value);
    }
  }

  void GetList(std::vector<TValue>& values) const {
    SHARED_ACQUIRE(TLock);
    values.clear();
    values.reserve(mMap.size());
    for (const auto& itr : mMap) values.push_back(itr.second);
  }

  bool Enum(TValue* list, size_t* size) const {
    if (!list && !size) return false;
    SHARED_ACQUIRE(TLock);
    if (size) {
      *size = mMap.size() * sizeof(TValue);
      if (!list) return true;
    }
    for (auto& itr : mMap) {
      *list = itr.second;
      AdjustValue(*list);
      ++list;
    }
    return true;
  }

  bool GetInfo(const TKey& key, TValue* valuePtr) const {
    TValue value;
    if (!Get(key, value)) return false;
    if (valuePtr) *valuePtr = value;
    return true;
  }

  TMap& GetDataUnsafe() { return mMap; }

  virtual void AdjustValue(TValue& value) const = 0;

 protected:
  virtual const char* jsonKey() const = 0;
  virtual TKey makeKey(const TValue& value) const = 0;

 private:
  TMap mMap;

  bool addNoLock(const TValue& value) {
    mMap[makeKey(value)] = value;
    return true;
  }

  bool getWhere(TValuePred predicate, TValue* value) {
    SHARED_ACQUIRE(TLock);
    for (const auto& itr : mMap) {
      if (!predicate(itr.second)) continue;
      if (value) *value = itr.second;
      return true;
    }
    return false;
  }
};

template <SectionLock TLock, class TKey, class TValue, class TSerializer,
          class THash = std::hash<TKey>>
using SerializableUnorderedMap =
    SerializableTMap<TLock, TKey, TValue,
                     std::unordered_map<TKey, TValue, THash>, TSerializer>;

template <SectionLock TLock, class TKey, class TValue, class TSerializer,
          class TCompare = std::less<TKey>>
using SerializableMap =
    SerializableTMap<TLock, TKey, TValue, std::map<TKey, TValue, TCompare>,
                     TSerializer>;

template <SectionLock TLock, class TValue, class TSerializer>
struct SerializableModuleRangeMap
    : SerializableMap<TLock, ModuleRange, TValue, TSerializer,
                      ModuleRangeCompare> {
  static ModuleRange VaKey(duint start, duint end) {
    auto moduleBase = ModBaseFromAddr(start);
    return ModuleRange(ModHashFromAddr(moduleBase),
                       Range(start - moduleBase, end - moduleBase));
  }
};

template <SectionLock TLock, class TValue, class TSerializer>
struct SerializableModuleHashMap
    : SerializableUnorderedMap<TLock, duint, TValue, TSerializer> {
  static duint VaKey(duint addr) { return ModHashFromAddr(addr); }

  void DeleteRangeWhere(
      duint start, duint end,
      std::function<bool(duint, duint, const TValue&)> inRange) {
    // Are all comments going to be deleted?
    // 0x00000000 - 0xFFFFFFFF
    if (start == 0 && end == ~0) {
      this->Clear();
    } else {
      // Make sure 'Start' and 'End' reference the same module
      duint moduleBase = ModBaseFromAddr(start);

      if (moduleBase != ModBaseFromAddr(end)) return;

      // Virtual -> relative offset
      start -= moduleBase;
      end -= moduleBase;

      this->DeleteWhere([start, end, inRange](const TValue& value) {
        return inRange(start, end, value);
      });
    }
  }
};

struct AddrInfo {
  duint modhash;
  duint addr;
  bool manual;

  std::string mod() const { return ModNameFromHash(modhash); }
};

template <class TValue>
struct AddrInfoSerializer : JSONWrapper<TValue> {
  static_assert(std::is_base_of<AddrInfo, TValue>::value,
                "TValue is not derived from AddrInfo");

  bool Save(const TValue& value) override {
    this->setString("module", value.mod());
    this->setHex("address", value.addr);
    this->setBool("manual", value.manual);
    return true;
  }

  bool Load(TValue& value) override {
    value.manual = true;  // legacy support
    this->getBool("manual", value.manual);
    std::string mod;
    if (!this->getString("module", mod)) return false;
    value.modhash = ModHashFromName(mod.c_str());
    return this->getHex("address", value.addr);
  }
};

template <SectionLock TLock, class TValue, class TSerializer>
struct AddrInfoHashMap : SerializableModuleHashMap<TLock, TValue, TSerializer> {
  static_assert(std::is_base_of<AddrInfo, TValue>::value,
                "TValue is not derived from AddrInfo");
  static_assert(std::is_base_of<AddrInfoSerializer<TValue>, TSerializer>::value,
                "TSerializer is not derived from AddrInfoSerializer");

  void AdjustValue(TValue& value) const override {
    value.addr += ModBaseFromName(value.mod().c_str());
  }

  bool PrepareValue(TValue& value, duint addr, bool manual) {
    if (!MemIsValidReadPtr(addr)) return false;
    auto base = ModBaseFromAddr(addr);
    value.modhash = ModHashFromAddr(base);
    value.manual = manual;
    value.addr = addr - base;
    return true;
  }

  void DeleteRange(duint start, duint end, bool manual) {
    this->DeleteRangeWhere(
        start, end, [manual](duint start, duint end, const TValue& value) {
          if (manual ? !value.manual
                     : value.manual)  // ignore non-matching entries
            return false;
          return value.addr >= start && value.addr < end;
        });
  }

 protected:
  duint makeKey(const TValue& value) const override {
    return value.modhash + value.addr;
  }
};

#endif  // _SERIALIZABLEMAP_H